'''
本次代码中，相较之前，进行了如下优化：
1. 不再对每个状态下的奖励做计算，因为对于其他步数较长的游戏来说，效率底下
2. 使用熵对计算出来的结果值进行惩罚，防止出现概率单一的动作分布
3. 采用N步展开的方式计算（todo 如何实现）
4. todo 是否有采用p183 9.3.4的样本之间的相关性
'''
#!/usr/bin/env python3
import gymnasium as gym
import ptan
import numpy as np
from tensorboardX import SummaryWriter

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

GAMMA = 0.99
LEARNING_RATE = 0.001
ENTROPY_BETA = 0.01 # 熵奖励的规模，todo 用于控制惩罚的力度
BATCH_SIZE = 8 # 每次训练的样本数，由于样本需要足够新，所以这里样本数不能太大

REWARD_STEPS = 10 # 指定bellman向前展开的步数，估算每次转换的总奖励损失，这里会导致
# 每次采样游戏进行的步数最大为为REWARD_STEPS个，但返回的样本仅包含第一个状态、以及该状态下指定的动作、最后一步的状态、以及所有步数的总奖励


class PGN(nn.Module):
    def __init__(self, input_size, n_actions):
        '''
        环境的观测维度，以及动作的维度
        '''
        super(PGN, self).__init__()

        self.net = nn.Sequential(
            nn.Linear(input_size, 128),
            nn.ReLU(),
            nn.Linear(128, n_actions)
        )

    def forward(self, x):
        return self.net(x)


if __name__ == "__main__":
    env = gym.make("CartPole-v0")
    writer = SummaryWriter(comment="-cartpole-pg")

    net = PGN(env.observation_space.shape[0], env.action_space.n)
    print(net)

    # 构建策略梯度网络代理，这里的apply_softmax=True，表示输出的是softmax之后的概率分布
    agent = ptan.agent.PolicyAgent(net, preprocessor=ptan.agent.float32_preprocessor,
                                   apply_softmax=True)
    # 与之前的对比，这里多了展开步数的配置，通过展开步数，每次迭代采样的最大样本数为steps_count个游戏步骤
    # GAMMA用于计算每次采样中每步下获取的总奖励
    exp_source = ptan.experience.ExperienceSourceFirstLast(env, agent, gamma=GAMMA, steps_count=REWARD_STEPS)

    optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE)

    total_rewards = []
    step_rewards = []
    step_idx = 0
    done_episodes = 0
    reward_sum = 0.0

    batch_states, batch_actions, batch_scales = [], [], []

    for step_idx, exp in enumerate(exp_source):
        # 根据奖励的总和，计算策略规模的基线 p182 9.3.2
        # 这里因为Pong游戏一轮游戏有成百上千帧，所以不能等待游戏结束
        # 而采用类似N步展开的方式，每次采样的样本数为REWARD_STEPS个游戏步骤
        # 因为不能等待游戏结束，所以这里采用计算回报基线的方式去衡量每个状态下的奖励
        # 这里基线是累计奖励的平均值
        reward_sum += exp.reward
        baseline = reward_sum / (step_idx + 1)
        writer.add_scalar("baseline", baseline, step_idx)
        batch_states.append(exp.state)
        batch_actions.append(int(exp.action))
        # 使用奖励基线是为了策略梯度下，PG具备很大的方差，这里使用奖励基线进行惩罚
        # 避免游戏采用绕路的策略，比如明明可以走几步就完成目标，游戏结束，但是采用绕路的方式完成游戏
        # 如果某个奖励过大，也会导致智能体更加倾向于采用这种方式，而不进行探索
        batch_scales.append(exp.reward - baseline)

        # 去除了计算每个状态下的总奖励

        # 获取最近一次游戏运行的到结束时的激励状态
        new_rewards = exp_source.pop_total_rewards()
        if new_rewards:
            # 将最后一次游戏完整运行的情况打印出来，如果激励平均值达到了指定目标，
            # 则判定为训练成功，退出循环
            done_episodes += 1
            reward = new_rewards[0]
            total_rewards.append(reward)
            mean_rewards = float(np.mean(total_rewards[-100:]))
            print("%d: reward: %6.2f, mean_100: %6.2f, episodes: %d" % (
                step_idx, reward, mean_rewards, done_episodes))
            writer.add_scalar("reward", reward, step_idx)
            writer.add_scalar("reward_100", mean_rewards, step_idx)
            writer.add_scalar("episodes", done_episodes, step_idx)
            if mean_rewards > 195:
                print("Solved in %d steps and %d episodes!" % (step_idx, done_episodes))
                break

        if len(batch_states) < BATCH_SIZE:
            continue

        states_v = torch.FloatTensor(batch_states)
        batch_actions_t = torch.LongTensor(batch_actions)
        batch_scale_v = torch.FloatTensor(batch_scales)

        optimizer.zero_grad()
        logits_v = net(states_v)
        # 这里为什么要采用log_softmax
        # 因为策略梯度实际上就是直接输出指定每一个动作的概率，所以这里采用log_softmax
        log_prob_v = F.log_softmax(logits_v, dim=1)
        # 根据计算出来的执行动作概率，使用batch_actions来从中选择实际执行的动作对应的概率
        # batch_scale_v是奖励基线，防止激励差异过大导致训练出现偏差
        # todo batch_scale_v实际的原理
        # 这里原理和cartpole_reinforce一样，计算出每个动作的概率，用实际执行的动作选择的动作概率
        # 乘以奖励基线，得到的就是每个动作的概率乘以奖励基线，为了使最终值越大，会加大奖励基线为正的动作
        # 减小奖励基线为负的动作
        log_prob_actions_v = batch_scale_v * log_prob_v[range(BATCH_SIZE), batch_actions_t]
        loss_policy_v = -log_prob_actions_v.mean()

        # 这边是通过计算熵，防止智能体在初期时没有进行探索 p183 9.3.3
        # 书上说通过信息熵，防止智能体得到过于肯定的策略
        prob_v = F.softmax(logits_v, dim=1)
        # H(p) = - Σ p(x) log p(x) 信息熵的计算公式，这里就是计算信息熵
        # 这里对应了书中公式
        entropy_v = -(prob_v * log_prob_v).sum(dim=1).mean()
        # 这里添加负号，是为了在最小化损失函数时，使得信息熵越大越好
        # 因为为了使损失最小，那么entropy_v就要越大，所以添加负号，鼓励智能体进行探索
        entropy_loss_v = -ENTROPY_BETA * entropy_v
        loss_v = loss_policy_v + entropy_loss_v

        loss_v.backward()
        optimizer.step()

        # calc KL-div 计算kl散度
        # 用来比较优化前与优化后的策略差异
        # 这里比较的是优化前与优化后的得到的概率分布的差异
        new_logits_v = net(states_v)
        new_prob_v = F.softmax(new_logits_v, dim=1)
        kl_div_v = -((new_prob_v / prob_v).log() * prob_v).sum(dim=1).mean()
        writer.add_scalar("kl", kl_div_v.item(), step_idx)

        # 计算梯度和l2范数图
        # 记录一个网络中所有参数的梯度的最大值，平均值，以及参数的l2范数
        # 了解训练动态
        grad_max = 0.0
        grad_means = 0.0
        grad_count = 0
        for p in net.parameters():
            grad_max = max(grad_max, p.grad.abs().max().item())
            grad_means += (p.grad ** 2).mean().sqrt().item()
            grad_count += 1

        writer.add_scalar("baseline", baseline, step_idx)
        writer.add_scalar("entropy", entropy_v.item(), step_idx)
        writer.add_scalar("batch_scales", np.mean(batch_scales), step_idx)
        writer.add_scalar("loss_entropy", entropy_loss_v.item(), step_idx)
        writer.add_scalar("loss_policy", loss_policy_v.item(), step_idx)
        writer.add_scalar("loss_total", loss_v.item(), step_idx)
        writer.add_scalar("grad_l2", grad_means / grad_count, step_idx)
        writer.add_scalar("grad_max", grad_max, step_idx)

        batch_states.clear()
        batch_actions.clear()
        batch_scales.clear()

    writer.close()
